import numpy as np
from pyswarm import pso

###############################################################################

print('*' * 65)
print('Example minimization of 4th-order banana function (no constraints)')


def myfunc(x):
    x1 = x[0]
    x2 = x[1]
    return x1 ** 4 - 2 * x2 * x1 ** 2 + x2 ** 2 + x1 ** 2 - 2 * x1 + 5


lb = [-3, -1]
ub = [2, 6]

xopt1, fopt1, pxopt1,pfopt1= pso(myfunc, lb, ub, particle_output=True)

print('The optimum is at:')
print('    {}'.format(xopt1))
print('Optimal function value:')
print('    myfunc: {}'.format(fopt1))
print('The optimum of per particle is at:')
print('    {}'.format(pxopt1))
print('Optimal function value of per particle:')
print('    myfunc: {}'.format(pfopt1))

###############################################################################

print('*' * 65)
print('Example minimization of 4th-order banana function (with constraint)')


def mycon(x):
    x1 = x[0]
    x2 = x[1]
    return [-(x1 + 0.25) ** 2 + 0.75 * x2]


xopt2, fopt2 = pso(myfunc, lb, ub, f_ieqcons=mycon)

print('The optimum is at:')
print('    {}'.format(xopt2))
print('Optimal function value:')
print('    myfunc: {}'.format(fopt2))
print('    mycon : {}'.format(mycon(xopt2)))

###############################################################################

print('*' * 65)
print('Engineering example: minimization of twobar truss weight, subject to')
print('  Yield Stress <= 100 kpsi')
print('  Yield Stress <= Buckling Stress')
print('  Deflection   <= 0.25 inches')


def weight(x, *args):
    H, d, t = x  # all in inches
    B, rho, E, P = args
    return rho * 2 * np.pi * d * t * np.sqrt((B / 2) ** 2 + H ** 2)


def stress(x, *args):
    H, d, t = x  # all in inches
    B, rho, E, P = args
    return (P * np.sqrt((B / 2) ** 2 + H ** 2)) / (2 * t * np.pi * d * H)


def buckling_stress(x, *args):
    H, d, t = x  # all in inches
    B, rho, E, P = args
    return (np.pi ** 2 * E * (d ** 2 + t ** 2)) / (8 * ((B / 2) ** 2 + H ** 2))


def deflection(x, *args):
    H, d, t = x  # all in inches
    B, rho, E, P = args
    return (P * np.sqrt((B / 2) ** 2 + H ** 2) ** 3) / (2 * t * np.pi * d * H ** 2 * E)


def mycons(x, *args):
    strs = stress(x, *args)
    buck = buckling_stress(x, *args)
    defl = deflection(x, *args)
    return [100 - strs, buck - strs, 0.25 - defl]


B = 60  # inches
rho = 0.3  # lb/in^3
E = 30000  # kpsi
P = 66  # lb (force)
args = (B, rho, E, P)
lb = [10, 1, 0.01]
ub = [30, 3, 0.25]
xopt4, fopt4 = pso(weight, lb, ub, f_ieqcons=mycons, args=args)

print('The optimum is at:')
print('    {}'.format(xopt4))
print('Optimal function values:')
print('    weight         : {}'.format(fopt4))
print('Constraint functions:')
print('    stress         : {}'.format(stress(xopt4, *args)))
print('    buckling stress: {}'.format(buckling_stress(xopt4, *args)))
print('    deflection     : {}'.format(deflection(xopt4, *args)))
